from torchvision.datasets import CIFAR10
from typing import Optional, Callable, Tuple, Any
from PIL import Image


class ShiftCIFAR10(CIFAR10):
    def __init__(
        self,
        root: str,
        train: bool = True,
        transform: Optional[Callable] = None,
        target_transform: Optional[Callable] = None,
        download: bool = False,
        num_classes: int = 10  # 新增参数，用于标签循环
    ) -> None:
        super().__init__(root, train, transform, target_transform, download)
        self.num_classes = num_classes

    def __getitem__(self, index: int) -> Tuple[Any, Any]:
        img, target = super().__getitem__(index)
        # 修改标签
        target = (target + 1) % self.num_classes
        return img, target
  